#pragma once
#include "matrix.h"
#include <string>
#include <vector>
#include <unordered_map>
#include <memory>
#include <fstream>
#include <sstream>
#include <algorithm>
#include <cctype>

class Part;

/**
 * @brief Represents an instance of a part in the assembly
 * 
 * An instance applies transformations (translation and rotation) to a base part
 */
class Instance {
private:
    std::shared_ptr<Part> ref_part;
    std::string name;
    std::vector<float> translate;
    std::vector<std::vector<float>> rotation; // [axis_point1, axis_point2, angle]

public:
    Instance(std::shared_ptr<Part> part, const std::string& inst_name = "")
        : ref_part(part), name(inst_name), translate(3, 0.0f) {}
    
    // Getters
    const std::string& getName() const { return name; }
    std::shared_ptr<Part> getRefPart() const { return ref_part; }
    const std::vector<float>& getTranslation() const { return translate; }
    const std::vector<std::vector<float>>& getRotation() const { return rotation; }
    
    // Setters
    void setTranslation(const std::vector<float>& trans) { translate = trans; }
    void setRotation(const std::vector<std::vector<float>>& rot) { rotation = rot; }
    
    /**
     * @brief Get node coordinates after applying instance transformations
     * @param nodeID Node identifier
     * @return Transformed coordinates [x, y, z]
     */
    std::vector<float> getNodeCoord(int nodeID);
    
    /**
     * @brief Get element information from the referenced part
     * @param elemID Element identifier
     * @return Pair of element type and node indices
     */
    std::pair<std::string, std::vector<int>> getElemInfo(int elemID);
};

/**
 * @brief Represents a part with nodes and elements
 */
class Part {
private:
    std::string name;
    Matrix<float> node_array;
    std::unordered_map<std::string, Matrix<int>> elem_types;
    std::unordered_map<int, int> node_id_to_index;
    std::unordered_map<int, std::pair<std::string, int>> elem_info;
    std::unordered_map<std::string, int> elem_type_offsets;

public:
    Part(const std::string& part_name = "") : name(part_name) {}
    
    // Getters
    const std::string& getName() const { return name; }
    const Matrix<float>& getNodeArray() const { return node_array; }
    const std::unordered_map<std::string, Matrix<int>>& getElemTypes() const { return elem_types; }
    const std::unordered_map<int, int>& getNodeIdToIndex() const { return node_id_to_index; }
    const std::unordered_map<int, std::pair<std::string, int>>& getElemInfoMap() const { return elem_info; }
    
    // Setters
    void setNodeArray(const Matrix<float>& nodes) { node_array = nodes; }
    void addElemType(const std::string& type, const Matrix<int>& elements) { elem_types[type] = elements; }
    void addNodeMapping(int node_id, int index) { node_id_to_index[node_id] = index; }
    void addElemInfo(int elem_id, const std::string& type, int index) { elem_info[elem_id] = std::make_pair(type, index); }
    
    /**
     * @brief Calculate element type offsets for sequential indexing
     */
    void elemSummary();
    
    /**
     * @brief Get sequential index of an element across all types
     * @param elem_id Element identifier
     * @return Sequential index or -1 if not found
     */
    int getElemSeqIndex(int elem_id);
    
    /**
     * @brief Get total number of elements in the part
     * @return Total element count
     */
    int getElemTotalNum() const;
    
    /**
     * @brief Get node coordinates from the part
     * @param nodeID Node identifier
     * @return Coordinates [x, y, z]
     */
    std::vector<float> getNodeCoord(int nodeID);
    
    /**
     * @brief Get element information
     * @param elemID Element identifier
     * @return Pair of element type and node indices
     */
    std::pair<std::string, std::vector<int>> getElemInfo(int elemID);
};

/**
 * @brief Represents the assembly containing parts and instances
 */
class Assembly {
private:
    std::unordered_map<std::string, std::shared_ptr<Instance>> instances;
    std::unordered_map<std::string, std::shared_ptr<Part>> parts;

public:
    // Instance management
    void addInstance(std::shared_ptr<Instance> instance);
    std::shared_ptr<Instance> getInstance(const std::string& name) const;
    std::unordered_map<std::string, std::shared_ptr<Instance>> getInstances() const { return instances; }
    std::vector<std::string> getInstanceList() const;
    
    // Part management
    void addPart(const std::string& name, std::shared_ptr<Part> part);
    std::shared_ptr<Part> getPart(const std::string& name) const;
};

/**
 * @brief Parser for Abaqus .inp files
 */
class AbaqusInpParser {
private:
    std::string filename;
    Assembly assembly;
    std::shared_ptr<Part> current_part;
    std::shared_ptr<Instance> current_instance;
    std::string current_elem_type;
    
    // Helper functions
    bool startsWith(const std::string& str, const std::string& prefix) const {
        return str.rfind(prefix, 0) == 0;
    }
    
    bool startsWithIgnoreCase(const std::string& str, const std::string& prefix) const {
        if (str.length() < prefix.length()) return false;
        std::string str_prefix = str.substr(0, prefix.length());
        std::string prefix_lower = prefix;
        std::transform(str_prefix.begin(), str_prefix.end(), str_prefix.begin(), ::tolower);
        std::transform(prefix_lower.begin(), prefix_lower.end(), prefix_lower.begin(), ::tolower);
        return str_prefix == prefix_lower;
    }

public:
    AbaqusInpParser(const std::string& file) : filename(file) {}
    
    /**
     * @brief Parse the input file and build assembly structure
     * @return Assembly object containing parts and instances
     */
    Assembly parse();

private:
    // Block parsing methods
    void parseNodeBlock(const std::vector<std::string>& lines, size_t& i);
    void parseElementBlock(const std::vector<std::string>& lines, size_t& i);
    std::string parseElementType(const std::string& line);
    std::vector<float> parseTranslate(const std::string& line);
    std::vector<std::vector<float>> parseRotate(const std::string& line);
};